from langchain_core.documents import Document
from langchain_core.messages import AnyMessage
from langgraph.graph.message import MessagesState, add_messages
from typing_extensions import List, TypedDict, Annotated

from HomeBuddyAgent.utils.state import reduce_device_results, reduce_device_configs, reduce_device_calls
from common.structs import DeviceModelFactory, DeviceCalls, ConfigT, DeviceResult


class State(MessagesState):
    question: str
    device_configs: Annotated[List[Document], reduce_device_configs]
    answer: str
    feed_back: bool
    factory: DeviceModelFactory
    device_calls: Annotated[DeviceCalls[ConfigT], reduce_device_calls]
    device_call_results: Annotated[List[DeviceResult], reduce_device_results]


class StateInput(TypedDict):
    question: str


class StateOutput(TypedDict):
    answer: str
    messages: Annotated[list[AnyMessage], add_messages]
